(** A function [f] decides a predicate [p] if it returns [true] exactly
    when [p] holds and [false] exactly when [p] does not hold. *)
predicate decide (f, p) = 
    (forall x. (post (f) (x) (true)  <=> p (x)))
and (forall x. (post (f) (x) (false) <=> not (p (x))))
and (forall x. pre (f) (x)) 

(** The following will be imported from the standard library in the future. *)

(* [&] decides the boolean conjunction. *)
deferred val & : (bool * bool) -> bool
where (forall (x, y). pre (&) (x, y))
  and (forall (x, y). post (&) (x, y) (true) <=> (x = true and y = true))

(* [+'] computes the addition between two integers. *)
deferred val +? : (int * int) -> int
where (forall (x, y). pre (+?) (x, y))
  and (forall (x, y, z). post (+?) (x, y) (z) <=> z = x + y)

(* [-'] computes the substraction between two integers. *)
deferred val -? : (int * int) -> int
where (forall (x, y). pre (-?) (x, y))
  and (forall (x, y, z). post (-?) (x, y) (z) <=> z = x - y)

(* [max] is the logical function that returns the maximum of two integers. *)
deferred logic max  : (int * int) ~> int
where forall (x, y). (x < y => max (x, y) = y) and (y <= x => max (x, y) = x)

(* [max'] is the implementation of the [max] logical function. *)
deferred val max? : (int * int) -> int
where (forall (x, y). pre (max?) (x, y))
  and (forall (x, y, z). post (max?) (x, y) (z) <=> z = max (x, y))

(* The distance between two integers. *)
deferred logic dist : (int * int) ~> int
where forall (x, y). 
  ((x > y) => dist (x, y) = (x - y)) and 
  ((not (x > y)) => dist (x, y) = (y - x))

(* Comparison operators over integers. *)
deferred val <=?  : (int * int) -> bool where decide (<=?, <=)
deferred val >=?  : (int * int) -> bool where decide (>=?, >=)
deferred val >?   : (int * int) -> bool where decide (>?, >)
deferred val =?   : (int * int) -> bool
where (forall (x, y). pre ( =? ) (x, y))
  and (forall (x, y). post (=?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=?) (x, y) (false) <=> not (x = y))

(** Elements of the sets. *)

(* The type of the elements. *)
deferred type elt : *

(* Their equality is decidable. *)
deferred val =elt?  : (elt * elt) -> bool
where (forall x. pre (=elt?) (x))
  and (forall (x, y). post (=elt?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=elt?) (x, y) (false) <=> not (x = y))

(* There exists an order. *)
deferred predicate <elt : (elt * elt) ~> prop 
where (forall x. not (x <elt x))
  and (forall (x, y). (not (x <elt y) and not (x = y)) => y <elt x)
  and (forall (x, y, z). (x <elt y and y <elt z) => x <elt z)
  and (forall (x, y). (x <elt y) or (y <elt x) or (x = y))

(* This order is decidable. *)
deferred val <elt? : (elt * elt) -> bool where decide (<elt?, <elt)

(** The sets. *)

(* The type of an abstract set. *)
deferred type set : *

(* The membership predicate. *)
deferred predicate %in : (elt * set) ~> prop

(* The empty set. *)
deferred val empty : set
where forall x [x]. not (x %in empty)

deferred logic singleton : elt ~> set
where forall (x, y). x %in (singleton (y)) <=> (x = y)

deferred logic @ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 @ s2) <=> ((x %in s1) or (x %in s2))
deferred logic ^ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 ^ s2) <=> ((x %in s1) and (x %in s2))

deferred logic \\ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 \\ s2) <=> ((x %in s1) and not (x %in s2))

deferred predicate %C : (set * set) ~> prop
where forall (s1, s2). (s1 %C s2) <=> (forall x. (x %in s1) => (x %in s2))

lemma flatten_union: forall (x, s1, s2, s3). 
  x %in (s1 @ s2 @ s3) <=> (x %in s1 or x %in s2 or x %in s3)

lemma transitive_subset: forall (s1, s2, s3).
  (s1 %C s2 and s2 %C s3) => s1 %C s3

(* Inferior bound of a set. *)
predicate inf (x, s) = forall y. y %in s => x <elt y

lemma: forall (x, s1, s2). (inf (x, s1) and s2 %C s1) => inf (x, s2)
lemma: forall (x, s1, s2). (inf (x, s1 @ s2)) <=> (inf (x, s1) and inf (x, s2))
lemma: forall (s, x, y). (inf (x, s) and y <elt x) => inf (y, s)
lemma: forall (x, y). (inf (x, singleton (y)) <=> x <elt y)

(* Superior bound of a set. *)
predicate sup (x, s) = forall y. y %in s => y <elt x

lemma: forall (x, s1, s2). (sup (x, s1) and s2 %C s1) => sup (x, s2)
lemma: forall (x, s1, s2). (sup (x, s1 @ s2)) <=> (sup (x, s1) and sup (x, s2))
lemma: forall (s, x, y). (sup (x, s) and x <elt y) => sup (y, s)
lemma: forall (x, y). (sup (x, singleton (y)) <=> y <elt x)

(* Extensional equality. *)
predicate === (s1, s2) = forall x. (x %in s1) <=> (x %in s2)

(* Extensionality axiom. *)
axiom: forall (s1, s2). (s1 === s2) <=> (s1 = s2) 

(** Trees. *)

(* Inductive type for tree. *)
type tree =
  | Empty : tree
  | Node  : (int * tree * elt * tree) -> tree

(* The abstraction of trees to sets. *)
logic elements (t) = 
  | Empty -> empty
  | Node (h, l, x, r) -> elements (l) @ singleton (x) @ elements (r)

lemma root_in_elements: forall (h, l, x, r). x %in elements (Node (h, l, x, r))
lemma subl_in_elements: forall (h, l, x, r). 
  elements (l) %C elements (Node (h, l, x, r))
lemma subr_in_elements: forall (h, l, x, r). 
  elements (r) %C elements (Node (h, l, x, r))
lemma subl2_in_elements: forall (h, hl, ll, lx, lr, x, r). 
    elements (ll) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))
and elements (lr) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))
lemma subr2_in_elements: forall (h, l, x, hr, rl, rx, rr). 
  elements (rl) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))
and elements (rr) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))

inductive bst : tree ~> prop =
  | bst (Empty)
  | forall (h, l, x, r). 
      bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))
      => bst (Node (h, l, x, r))

predicate bst_node (l, x, r) = 
   bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))


type iterator =
  | Nil  : iterator
  | Cons : (tree * iterator) -> iterator

logic remaining (i) =
  | Nil -> empty
  | Cons (t, ts) -> elements (t) @ remaining (ts)

type next_result = 
  | None : next_result
  | Some : (iterator * elt) -> next_result

inductive ok : iterator ~> prop = 
| ok (Nil)
| forall (t, ts). (elements (t) ^ remaining (ts)) === empty and bst (t) and ok (ts) => ok (Cons (t, ts))

let iterator (t) where bst (t) 
returns i where (ok (i) and remaining (i) === elements (t)) =
  Cons (t, Nil)

let rec next (i) where ok (i) 
returns oix 
where ((oix = None => remaining (i) === empty)
and (forall (i', x). oix = Some (i', x) => (remaining (i) === (singleton (x) @ remaining (i'))
                                      and not (x %in remaining (i')) and ok (i')))) =
  match i with
  | Nil -> None
  | Cons (Empty, ts) -> next (ts)
  | Cons (Node (h, l, x, r), ts) -> Some (Cons (l, Cons(r, ts)), x)
  end

deferred logic cardinal : set ~> int
where cardinal (empty) = 0
  and (forall y. cardinal (singleton (y)) = 1)
  and (forall (x, s). (not (x %in s)) => cardinal (singleton (x) @ s) = 1 + cardinal (s))

let cardinal? (t) where bst (t) returns n where (n = cardinal (elements (t))) =
  let rec count (i, n) where ok (i) and n + cardinal (remaining (i)) = cardinal (elements (t))
  returns n' where (n' = cardinal (elements (t))) =
    match next (i) with
    | None -> n
    | Some (i', x) -> count (i', n +? 1)
    end
  in
    count (iterator (t), 0)
